import asyncio
import logging

import tornado
import yaml

from algorithms.SAM.build_sam import sam_model_registry
from algorithms.SAM.predictor import SamPredictor
from server.anno_handler import (FindModel,
                                 init_pg_connection_pool,
                                 CreateWorkspace,
                                 FindAnnoDatas,
                                 SaveLabel,
                                 FindLabels,
                                 DeleteLabels,
                                 FindImage,
                                 SaveLabelInstance,
                                 FindLabelInstances,
                                 DeleteLabelInstances,
                                 FindWorkspaces,
                                 SamPredict,
                                 CombineLabelInstances,
                                 FindLabelFiles, GenerateLabelFile)


def make_app(_config):
    init_pg_connection_pool(_config)
    _app = tornado.web.Application([
        (r"/elearning/findModel", FindModel),
        (r"/elearning/findWorkspaces", FindWorkspaces),
        (r"/elearning/createWorkspace", CreateWorkspace),
        (r"/elearning/findAnnoDatas", FindAnnoDatas),
        (r"/elearning/saveLabel", SaveLabel),
        (r"/elearning/findLabels", FindLabels),
        (r"/elearning/deleteLabels", DeleteLabels),
        (r"/elearning/([^/]+)?/([^/]+)?", FindImage),
        (r"/elearning/saveLabelInstance", SaveLabelInstance),
        (r"/elearning/findLabelInstances", FindLabelInstances),
        (r"/elearning/deleteLabelInstances", DeleteLabelInstances),
        (r"/elearning/combineLabelInstances", CombineLabelInstances),
        (r"/elearning/samPredict", SamPredict),
        (r"/elearning/findLabelFiles", FindLabelFiles),
        (r"/elearning/generateLabelFile", GenerateLabelFile),
    ],**_config)
    # 创建sam模型
    sam_config = _config["sam"]
    sam_model = sam_model_registry[sam_config["model_type"]](_config["app"]['model_path'],sam_config["checkpoint"])
    sam_model.to("cuda")
    _app.sam_predictor = SamPredictor(model=sam_model, device=_config["app"]['device'])
    return _app


async def stop_app(_app, _server):
    logging.warning("stopping elearning......")
    await _app.predictor.stop()
    _server.stop()
    await asyncio.sleep(1)
    logging.info("elearning is stopped")


def load_config():
    config_path = "config/application_linux.yaml"
    logging.info("loading config {}".format(config_path))
    with open(config_path, 'r', encoding='utf-8') as f:
        _config = yaml.safe_load(f)
    return _config


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO,
                        format='%(asctime)s.%(msecs)03d [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s',
                        datefmt='## %Y-%m-%d %H:%M:%S')
    logging.info("loading config")
    config = load_config()
    logging.info("initializing elearning")
    app = make_app(config)
    logging.info("elearning is listening on port {}".format(config["app"]["port"]))
    server = app.listen(config["app"]["port"])
    loop = asyncio.get_event_loop()
    try:
        loop.run_forever()
    except KeyboardInterrupt:
        logging.warning("sam_predictor is cancelled")
        logging.info("elearning is stopped")
    finally:
        loop.stop()
        loop.close()
